#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#pragma pack(push,1)
typedef struct {
    uint8_t boot_flag; // 0x80: 活动分区, 0x00: 非活动
    uint8_t start_chs[3]; // 起始CHS地址，写0
    uint8_t part_type; // 分区类型，0x83: Linux
    uint8_t end_chs[3]; // 结束CHS地址，写0
    uint32_t start_lba; // 起始LBA
    uint32_t sectors; // 分区大小（单位：扇区）
} PartitionEntry;

typedef struct {
    uint8_t boot_code[446]; // Bootloader 代码区
    PartitionEntry parts[4]; // 最多4个主分区
    uint16_t signature; // 固定为0xAA55
} MBR;
#pragma pack(pop)

int main(int argc, char *argv[]) {
    if (argc < 4 || ((argc - 2) % 2 != 0) || ((argc - 2) / 2 > 4)) {
        printf("用法: %s <镜像文件> [起始1 大小1] [起始2 大小2] ... 最多4个分区\n", argv[0]);
        return 1;
    }

    const char *image_path = argv[1];
    int part_count = (argc - 2) / 2;

    MBR mbr;
    memset(&mbr, 0, sizeof(MBR));

    for (int i = 0; i < part_count; ++i) {
        uint32_t start = (uint32_t) atoi(argv[2 + i * 2]);
        uint32_t size = (uint32_t) atoi(argv[2 + i * 2 + 1]);

        mbr.parts[i].boot_flag = (i == 0) ? 0x80 : 0x00; // 第一个分区设为活动
        memset(mbr.parts[i].start_chs, 0, 3);
        mbr.parts[i].part_type = (i == 0) ? 0x83 : 0x0B;
        memset(mbr.parts[i].end_chs, 0, 3);
        mbr.parts[i].start_lba = start;
        mbr.parts[i].sectors = size;
    }

    mbr.signature = 0xAA55;

    FILE *fp = fopen(image_path, "r+b");
    if (!fp) {
        perror("打开镜像文件失败");
        return 1;
    }

    fseek(fp, 0, SEEK_SET);
    size_t written = fwrite(&mbr, 1, sizeof(MBR), fp);
    fclose(fp);

    if (written != sizeof(MBR)) {
        printf("写入MBR失败：仅写入了 %zu 字节\n", written);
        return 1;
    }

    printf("成功写入MBR，共 %d 个分区：\n", part_count);
    for (int i = 0; i < part_count; ++i) {
        printf("  分区%d：起始 %u，大小 %u 扇区%s\n",
               i + 1,
               mbr.parts[i].start_lba,
               mbr.parts[i].sectors,
               i == 0 ? "（活动）" : "");
    }

    return 0;
}
